{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "UNiYBVL25rQK"
   },
   "source": [
    "# Fine Tuning Transformer for Sentiment Classification"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Introduction\n",
    "\n",
    "In this tutorial we will be fine tuning a transformer model for the **Sentiment classification** problem. **Sentiment classification** is a special case of **Multiclass Classification**. In this case the classes represent the sentiment represented by the text.\n",
    "The number of classes are generally lesser than a standard multiclass classification proboem where the classes represent the polarity, in form of `postive`, `negative` and in some cases and additional `neutral` polarity.\n",
    "\n",
    "This is one of the most common business problems when trying to ascertain the sentiment of a statement made by your customer in a business setup.\n",
    "\n",
    "#### Flow of the notebook\n",
    "\n",
    "* As with all the tutorials previously, this notebook also follows a easy to follow steps. Making the process of fine tuning and training a Transformers model a straight forward task.\n",
    "* However, unlike the other notebooks, in the tutorial, most of the sections have been created into functions, and they are called from the `main()` in the end of the notebook. \n",
    "* This is done to leverage the [Weights and Biases Service](https://www.wandb.com/) WandB in short.\n",
    "* It is a experiment tracking, parameter optimization and artifact management service. That can be very easily integrated to any of the Deep learning or Machine learning frameworks. \n",
    "\n",
    "The notebook will be divided into separate sections to provide a organized walk through for the process used. This process can be modified for individual use cases. The sections are:\n",
    "\n",
    "1. [Preparing Environment and Importing Libraries](#section01)\n",
    "2. [Pre-Processing and Preparing the Dataset for data processing: Class](#section02)\n",
    "3. [Defining a Model/Network](#section07)\n",
    "4. [Fine Tuning the Model: Function](#section03)\n",
    "5. [Validating the Model Performance: Function](#section04)\n",
    "6. [Main Function](#section05)\n",
    "    * [Initializing WandB](#section501)\n",
    "    * [Importing and Pre-Processing the domain data](#section502)\n",
    "    * [Creation of Dataset and Dataloader](#section503)\n",
    "    * [Neural Network and Optimizer](#section504)\n",
    "    * [Training Model and Logging to WandB](#section505)\n",
    "    * [Validation and generation of Summary](#section506)\n",
    "\n",
    "\n",
    "#### Technical Details\n",
    "\n",
    "This script leverages on multiple tools designed by other teams. Details of the tools used below. Please ensure that these elements are present in your setup to successfully implement this script.\n",
    "\n",
    "- **Data**:\n",
    "\t- We are using the **IMDB Dataset** available at [Kaggle](https://www.kaggle.com/lakshmi25npathi/imdb-dataset-of-50k-movie-reviews)\n",
    "\t- This dataset is a collection of moview reviews obtained from IMDB website, the reviews are labled with a positive or negative sentiment. \n",
    "\t- There are approx. `50000` rows of data.  Where each row has the following data-point:\n",
    "\t\t- **review** : Review of a movie\n",
    "\t\t- **sentiment** : positive or negative\n",
    "\n",
    "\n",
    "- **Language Model Used**: \n",
    "    - This notebook uses ***RoBERTa*** as its base transformer model. [Research Paper](https://arxiv.org/abs/1907.11692)    \n",
    "    - ***RoBERTa*** was an incremental improvement in the ***BERT*** architecture with multiple tweaks in different domains.\n",
    "    - Some of the changes in RoBERTa were: Bigger training data, Dymanic Masking, Different Self Supervised training objective.\n",
    "    - You can have a detailed read of these changes at the following [link](https://medium.com/towards-artificial-intelligence/a-robustly-optimized-bert-pretraining-approach-f6b6e537e6a6). \n",
    "   - We will be leveraging on the ***RoBERTa*** implementation from the HuggingFace team.    \n",
    "   - [Documentation for python](https://huggingface.co/transformers/model_doc/t5.html)\n",
    "\n",
    "\n",
    "- **Hardware Requirements**: \n",
    "\t- Python 3.6 and above\n",
    "\t- Pytorch, Transformers and\n",
    "\t- All the stock Python ML Library\n",
    "\t- GPU/TPU enabled setup \n",
    "   \n",
    "\n",
    "- **Script Objective**:\n",
    "\t- The objective of this script is to fine tune ***RoBERTa*** to be able to classify wether the sentiment of a given text is positive or negative.\n",
    "\n",
    "---\n",
    "NOTE: \n",
    "We are using the Weights and Biases Tool-set in  this tutorial. The different components will be explained as we go through the article. This is an incremental work done in the summarization notebook.\n",
    "\n",
    "[Link](https://app.wandb.ai/abhimishra-91/transformers_tutorials_sentiment?workspace=user-abhimishra-91) to the Project on WandB"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='section01'></a>\n",
    "### Preparing Environment and Importing Libraries\n",
    "\n",
    "At this step we will be installing the necessary libraries followed by importing the libraries and modules needed to run our script. \n",
    "We will be installing:\n",
    "* transformers\n",
    "* wandb\n",
    "* packages to support tpu for pytorch\n",
    "\n",
    "Libraries imported are:\n",
    "* Pandas\n",
    "* Pytorch\n",
    "* Pytorch Utils for Dataset and Dataloader\n",
    "* Transformers\n",
    "* Roberta Model and Tokenizer\n",
    "* wandb\n",
    "\n",
    "Followed by that we will preapre the device to support TPU execution for training.\n",
    "\n",
    "Finally, we will be logging into the [wandb](https://www.wandb.com/) serice using the login command"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 844
    },
    "colab_type": "code",
    "id": "4k86Li6zL9Jz",
    "outputId": "6fb4449b-12cc-4917-84f0-6e17ff56b9e7"
   },
   "outputs": [],
   "source": [
    "# Installing NLP-Transformers library\n",
    "!pip install -q transformers\n",
    "\n",
    "# Installing wandb library for experiment tracking and hyper parameter optimization\n",
    "!pip install -q wandb\n",
    "\n",
    "# Code for TPU packages install\n",
    "!curl -q https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n",
    "!python pytorch-xla-env-setup.py --apt-packages libomp5 libopenblas-dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hCOWCYgQvgkP"
   },
   "outputs": [],
   "source": [
    "# Importing stock libraries\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader, RandomSampler, SequentialSampler\n",
    "\n",
    "# Importing lackages from our NLP-Hugging Package\n",
    "from transformers import RobertaConfig, RobertaModel, RobertaTokenizerFast, RobertaForSequenceClassification\n",
    "\n",
    "# Importing wand for logging and hyper-parameter tuning\n",
    "import wandb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "i7NMSGVKDwue"
   },
   "outputs": [],
   "source": [
    "# Setting up the accelerators\n",
    "\n",
    "# # GPU\n",
    "# from torch import cuda\n",
    "# device = 'cuda' if cuda.is_available() else 'cpu'\n",
    "\n",
    "# TPU\n",
    "import torch_xla\n",
    "import torch_xla.core.xla_model as xm\n",
    "device = xm.xla_device()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 87
    },
    "colab_type": "code",
    "id": "X9L1AxSQDugC",
    "outputId": "76ceca60-14c8-4baf-b410-ec2c5c179bc9"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: You can find your API key in your browser here: https://app.wandb.ai/authorize\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Paste an API key from your profile and hit enter: \n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Appending key for api.wandb.ai to your netrc file: /root/.netrc\n",
      "\u001b[32mSuccessfully logged in to Weights & Biases!\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# login to wandb\n",
    "!wandb login"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='section02'></a>\n",
    "### Pre-Processing and Preparing the Dataset for data processing: Class\n",
    "\n",
    "* We will start with creation of Preprocess class - This defines how the text is pre-processed before working on the tokenization, dataset and dataloader aspects of the workflow. In this class the dataframe is loaded and then the `sentiment` column is used to create a new column in the dataframe called `encoded_polarity` such that if:\n",
    "    * `sentiment = positive` then `encoded_polarity = 0`\n",
    "    * `sentiment = negative` then `encoded_polarity = 1`\n",
    "\n",
    "* Followed by this, the `sentiment` column is removed from the dataframe.\n",
    "* The `dataframe` and `encoded_polarity` dictionary are returned. \n",
    "* This method is called in the `run()` function.\n",
    "\n",
    "* After this we will work on the Dataset class - This defines how the text is pre-processed before sending it to the neural network. This dataset will be used the the Dataloader method that will feed  the data in batches to the neural network for suitable training and processing. \n",
    "* The Dataloader and Dataset will be used inside the `run()`.\n",
    "* Dataset and Dataloader are constructs of the PyTorch library for defining and controlling the data pre-processing and its passage to neural network. For further reading into Dataset and Dataloader read the [docs at PyTorch](https://pytorch.org/docs/stable/data.html)\n",
    "\n",
    "#### *CustomDataset* Dataset Class\n",
    "- This class is defined to accept the Dataframe as input and generate tokenized output that is used by the Roberta model for training. \n",
    "- We are using the Roberta tokenizer to tokenize the data in the `review` column of the dataframe. \n",
    "- The tokenizer uses the `encode_plus` method to perform tokenization and generate the necessary outputs, namely: `ids`, `attention_mask`\n",
    "- To read further into the tokenizer, [refer to this document](https://huggingface.co/transformers/model_doc/roberta.html#robertatokenizer)\n",
    "- `encoded_polarity` transformed into the `targets` tensor. \n",
    "- The *CustomDataset* class is used to create 2 datasets, for training and for validation.\n",
    "- *Training Dataset* is used to fine tune the model: **70% of the original data**\n",
    "- *Validation Dataset* is used to evaluate the performance of the model. The model has not seen this data during training. \n",
    "\n",
    "#### return_dataloader: Called inside the `run()`\n",
    "- `return_dataloader` function is used to for creating training and validation dataloader that load data to the neural network in a defined manner. This is needed because all the data from the dataset cannot be loaded to the memory at once, hence the amount of data loaded to the memory and then passed to the neural network needs to be controlled.\n",
    "- Internally the `return_dataloader` function calls the pytorch Dataloader class and the CustomDataset class to create the dataloaders for training and validation. \n",
    "- This control is achieved using the parameters such as `batch_size` and `max_len`.\n",
    "- Training and Validation dataloaders are used in the training and validation part of the flow respectively"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "yDe5ntrM9Kud"
   },
   "outputs": [],
   "source": [
    "# Preprocess class defines how the dataframe will be processed to generate  and removal of features based on thier usage.\n",
    "# A new encoded_polarity column is added that adds 0 and 1 to the column based on the positive and negative\n",
    "# The processing method will return both the dictionary, and the updated dataframe for further usage.\n",
    "\n",
    "class Preprocess:\n",
    "    def __init__(self, df):\n",
    "        \"\"\"\n",
    "        Constructor for the class\n",
    "        :param df: Input Dataframe to be pre-processed\n",
    "        \"\"\"\n",
    "        self.df = df\n",
    "        self.encoded_dict = dict()\n",
    "\n",
    "    def encoding(self, x):\n",
    "        if x not in self.encoded_dict.keys():\n",
    "            self.encoded_dict[x] = len(self.encoded_dict)\n",
    "        return self.encoded_dict[x]\n",
    "\n",
    "    def processing(self):\n",
    "        self.df['encoded_polarity'] = self.df['sentiment'].apply(lambda x: self.encoding(x))\n",
    "        self.df.drop(['sentiment'], axis=1, inplace=True)\n",
    "        return self.encoded_dict, self.df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "l00E-IDDdg0s"
   },
   "outputs": [],
   "source": [
    "# Creating a CustomDataset class that is used to read the updated dataframe and tokenize the text. \n",
    "# The class is used in the return_dataloader function\n",
    "\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, dataframe, tokenizer, max_len):\n",
    "        self.len = len(dataframe)\n",
    "        self.data = dataframe\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_len = max_len\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        text = str(self.data.review[index])\n",
    "        text = \" \".join(text.split())\n",
    "        inputs = self.tokenizer.encode_plus(\n",
    "            text,\n",
    "            None,\n",
    "            add_special_tokens=True,\n",
    "            max_length=self.max_len,\n",
    "            pad_to_max_length=True,\n",
    "            return_token_type_ids=True\n",
    "        )\n",
    "        ids = inputs['input_ids']\n",
    "        mask = inputs['attention_mask']\n",
    "\n",
    "        return {\n",
    "            'ids': torch.tensor(ids, dtype=torch.long),\n",
    "            'mask': torch.tensor(mask, dtype=torch.long),\n",
    "            'targets': torch.tensor(self.data.encoded_polarity[index], dtype=torch.float)\n",
    "        } \n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.len"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Oj5MC7xPhMoq"
   },
   "outputs": [],
   "source": [
    "# Creating a function that returns the dataloader based on the dataframe and the specified train and validation batch size. \n",
    "\n",
    "def return_dataloader(df, tokenizer, train_batch_size, validation_batch_size, MAX_LEN, train_size=0.7):\n",
    "    train_size = 0.7\n",
    "    train_dataset=df.sample(frac=train_size,random_state=200)\n",
    "    val_dataset=df.drop(train_dataset.index).reset_index(drop=True)\n",
    "    train_dataset = train_dataset.reset_index(drop=True)\n",
    "\n",
    "    print(\"FULL Dataset: {}\".format(df.shape))\n",
    "    print(\"TRAIN Dataset: {}\".format(train_dataset.shape))\n",
    "    print(\"VAL Dataset: {}\".format(val_dataset.shape))\n",
    "\n",
    "    training_set = CustomDataset(train_dataset, tokenizer, MAX_LEN)\n",
    "    validation_set = CustomDataset(val_dataset, tokenizer, MAX_LEN)\n",
    "\n",
    "    train_params = {'batch_size': train_batch_size,\n",
    "                'shuffle': True,\n",
    "                'num_workers': 1\n",
    "                }\n",
    "\n",
    "    val_params = {'batch_size': validation_batch_size,\n",
    "                    'shuffle': True,\n",
    "                    'num_workers': 1\n",
    "                    }\n",
    "\n",
    "    training_loader = DataLoader(training_set, **train_params)\n",
    "    validation_loader = DataLoader(validation_set, **val_params)\n",
    "    \n",
    "    return training_loader, validation_loader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='section07'></a>\n",
    "### Defining a Model/Network\n",
    "\n",
    "#### Neural Network\n",
    " - We will be creating a neural network with the `ModelClass`. \n",
    " - This network will have the Roberta Language model and a few by a `dropout` and `Linear` layer to obtain the final outputs. \n",
    " - The data will be fed to the Roberta Language model as defined in the dataset. \n",
    " - Final layer outputs is what will be compared to the `encoded_polarity` to determine the accuracy of models prediction. \n",
    " - We will initiate an instance of the network called `model`. This instance will be used for training and then to save the final trained model for future inference. \n",
    " - The `return_model` function is used in the `run()` to instantiate the model and set it up for TPU execution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MxaCC_XGh0ev"
   },
   "outputs": [],
   "source": [
    "# Creating the customized model, by adding a drop out and a dense layer on top of roberta to get the final output for the model. \n",
    "\n",
    "class ModelClass(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(ModelClass, self).__init__()\n",
    "        self.model_layer = RobertaModel.from_pretrained(\"roberta-base\")\n",
    "        self.pre_classifier = torch.nn.Linear(768, 768)\n",
    "        self.dropout = torch.nn.Dropout(0.3)\n",
    "        self.classifier = torch.nn.Linear(768, 2)\n",
    "\n",
    "    def forward(self, input_ids, attention_mask):\n",
    "        output_1 = self.model_layer(input_ids=input_ids, attention_mask=attention_mask)\n",
    "        hidden_state = output_1[0]\n",
    "        pooler = hidden_state[:, 0]\n",
    "        pooler = self.pre_classifier(pooler)\n",
    "        pooler = torch.nn.ReLU()(pooler)\n",
    "        pooler = self.dropout(pooler)\n",
    "        output = self.classifier(pooler)\n",
    "        return output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tr7ueXG-B3Vz"
   },
   "outputs": [],
   "source": [
    "# Function to return model based on the defination of Model Class\n",
    "\n",
    "def return_model(device):\n",
    "    model = ModelClass()\n",
    "    model = model.to(device)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rbY5DHfbDSuZ"
   },
   "outputs": [],
   "source": [
    "# Function to calcuate the accuracy of the model\n",
    "\n",
    "def calcuate_accu(big_idx, targets):\n",
    "    n_correct = (big_idx==targets).sum().item()\n",
    "    return n_correct"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='section03'></a>\n",
    "### Fine Tuning the Model: Function\n",
    "\n",
    "Here we define a training function that trains the model on the training dataset created above, specified number of times (EPOCH), An epoch defines how many times the complete data will be passed through the network. \n",
    "\n",
    "This function is called in the `run()`\n",
    "\n",
    "Following events happen in this function to fine tune the neural network:\n",
    "- The `epoch`, `model`, `device` details, `testing_ dataloader`, `optimizer` and `loss_function` are passed to the `train ()` when its called from the `run()`\n",
    "- The dataloader passes data to the model based on the batch size.\n",
    "- The output from the neural network: `outputs` is compared to the `targets` tensor and loss is calcuated using `loss_function()`\n",
    "- Loss value is used to optimize the weights of the neurons in the network.\n",
    "- After every 100 steps the loss value and accuracy is logged in the wandb service. This log is then used to generate graphs for analysis. Such as [these](https://app.wandb.ai/abhimishra-91/transformers_tutorials_sentiment?workspace=user-abhimishra-91)\n",
    "- After every epoch the loss and accuracy value is printed in the console. Also, logged into the wandb service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7GE89LJnkid7"
   },
   "outputs": [],
   "source": [
    "# Function to fine tune the model based on the epochs, model, tokenizer and other arguments\n",
    "\n",
    "def train(epoch, model, device, training_loader, optimizer, loss_function):\n",
    "    n_correct = 0\n",
    "    nb_tr_examples, nb_tr_steps = 0, 0\n",
    "    tr_loss = 0\n",
    "    model.train()\n",
    "    for _,data in enumerate(training_loader, 0):\n",
    "        ids = data['ids'].to(device, dtype = torch.long)\n",
    "        mask = data['mask'].to(device, dtype = torch.long)\n",
    "        targets = data['targets'].to(device, dtype = torch.long)\n",
    "\n",
    "        outputs = model(ids, mask).squeeze()\n",
    "        optimizer.zero_grad()\n",
    "        loss = loss_function(outputs, targets)\n",
    "        tr_loss += loss.item()\n",
    "        big_val, big_idx = torch.max(outputs.data, dim=1)\n",
    "        n_correct += calcuate_accu(big_idx, targets)\n",
    "\n",
    "        nb_tr_steps += 1\n",
    "        nb_tr_examples+=targets.size(0)\n",
    "        \n",
    "        if _%100==0:\n",
    "            loss_step = tr_loss/nb_tr_steps\n",
    "            accu_step = (n_correct*100)/nb_tr_examples \n",
    "            wandb.log({\"Training Loss per 100 steps\": loss_step})\n",
    "            wandb.log({\"Training Accuracy per 100 steps\": accu_step})\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        \n",
    "        # # When using GPU or GPU\n",
    "        # optimizer.step()\n",
    "        \n",
    "        # When using TPU\n",
    "        xm.optimizer_step(optimizer)\n",
    "        xm.mark_step()\n",
    "\n",
    "    print(f'The Total Accuracy for Epoch {epoch}: {(n_correct*100)/nb_tr_examples}')\n",
    "    epoch_loss = tr_loss/nb_tr_steps\n",
    "    epoch_accu = (n_correct*100)/nb_tr_examples\n",
    "    wandb.log({\"Training Loss Epoch\": epoch_loss})\n",
    "    wandb.log({\"Training Accuracy Epoch\": epoch_accu})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='section04'></a>\n",
    "### Validating the Model Performance: Function\n",
    "\n",
    "During the validation stage we pass the unseen data(Validation Dataset), trained model, and device details to the function to perform the validation run. This step generates new encoded_sentiment value for dataset that it has not seen during the training session. \n",
    "\n",
    "This is then compared to the actual encoded_sentiment, to give us the Validation Accuracy and Loss.\n",
    "\n",
    "This function is called in the `run()`\n",
    "\n",
    "This unseen data is the 30% of `IMBD Dataset` which was seperated during the Dataset creation stage. \n",
    "During the validation stage the weights of the model are not updated. We use the generate method for generating new text for the summary. \n",
    "\n",
    "The generated validation accuracy and loss are logged to wandb for every 100th step and per epoch. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2_1Hw4k-0M4r"
   },
   "outputs": [],
   "source": [
    "# Function to run the validation dataloader to validate the performance of the fine tuned model. \n",
    "\n",
    "def valid(epoch, model, device, validation_loader, loss_function):\n",
    "    n_correct = 0; total = 0\n",
    "    nb_tr_examples, nb_tr_steps = 0, 0\n",
    "    tr_loss = 0\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for _,data in enumerate(validation_loader, 0):\n",
    "            ids = data['ids'].to(device, dtype = torch.long)\n",
    "            mask = data['mask'].to(device, dtype = torch.long)\n",
    "            targets = data['targets'].to(device, dtype = torch.long)\n",
    "\n",
    "            outputs = model(ids, mask).squeeze()\n",
    "            loss = loss_function(outputs, targets)\n",
    "            tr_loss += loss.item()\n",
    "            big_val, big_idx = torch.max(outputs.data, dim=1)\n",
    "            n_correct += calcuate_accu(big_idx, targets)\n",
    "\n",
    "            nb_tr_steps += 1\n",
    "            nb_tr_examples+=targets.size(0)\n",
    "            \n",
    "            if _%100==0:\n",
    "                loss_step = tr_loss/nb_tr_steps\n",
    "                accu_step = (n_correct*100)/nb_tr_examples \n",
    "                wandb.log({\"Validation Loss per 100 steps\": loss_step})\n",
    "                wandb.log({\"Validation Accuracy per 100 steps\": accu_step})\n",
    "        \n",
    "    epoch_loss = tr_loss/nb_tr_steps\n",
    "    epoch_accu = (n_correct*100)/nb_tr_examples\n",
    "    wandb.log({\"Validation Loss Epoch\": epoch_loss})\n",
    "    wandb.log({\"Validation Accuracy Epoch\": epoch_accu})\n",
    "    print(f'The Validation Accuracy: {(n_correct*100)/nb_tr_examples}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='section05'></a>\n",
    "### Run Function\n",
    "\n",
    "The `run()` as the name suggests is the central location to run all the functions/flows created above in the notebook. The following steps are executed in the `run()`:\n",
    "\n",
    "\n",
    "<a id='section501'></a>\n",
    "#### Initializing WandB \n",
    "\n",
    "* The `run()` begins with initializing WandB run under a specific project. This command initiates a new run for each execution of this command. \n",
    "\n",
    "* We have seend wandb in action in one of the previous notebooks. Leveraging this notebook to log some additional metrics. \n",
    "\n",
    "* This particular tutorial is logged in the project: **[transformers_tutorials_sentiment](https://app.wandb.ai/abhimishra-91/transformers_tutorials_sentiment?workspace=user-abhimishra-91)**\n",
    "\n",
    "**One of the dadshboard from the project**\n",
    "![](meta/wandb-sentiment.jpg)\n",
    "\n",
    "* Visit the project page to see the details of different runs and what information is logged by the service. \n",
    "\n",
    "* Following the initialization of the WandB service we define configuration parameters that will be used across the tutorial such as `batch_size`, `epoch`, `learning_rate` etc.\n",
    "\n",
    "* These parameters are also passed to the WandB config. The config construct with all the parameters can be optimized using the Sweep service from WandB. Currently, that is outof scope of this tutorial. \n",
    "\n",
    "\n",
    "<a id='section502'></a>\n",
    "#### Importing and Pre-Processing the domain data\n",
    "\n",
    "We will be working with the data and preparing it for fine tuning purposes. \n",
    "*Assuming that the `IMDB Dataset.csv` is already downloaded in your `data` folder*\n",
    "\n",
    "* The file is imported as a dataframe and give it the headers as per the documentation.\n",
    "* Cleaning the file to remove the unwanted columns.\n",
    "* All these steps are done using the `Preprocess Class` defined above\n",
    "* The final Dataframe will be something like this:\n",
    "\n",
    "|review|encoded_polarity|\n",
    "|--|--|\n",
    "|summary-1|0|\n",
    "|summary-2|1|\n",
    "|summary-3|1|\n",
    "\n",
    "\n",
    "<a id='section503'></a>\n",
    "#### Creation of Dataset and Dataloader\n",
    "\n",
    "* The updated dataframe is divided into 70-20 ratio for test and validation. \n",
    "* Both the data-frames are passed to the `CustomerDataset` class for tokenization of the review and its sentiment.\n",
    "* The tokenization is done using roberta tokenizer.\n",
    "* Train and Validation parameters are defined and passed to the `pytorch Dataloader contstruct` to create `train` and `validation` data loaders.\n",
    "* These dataloaders will be passed to `train()` and `validate()` respectively for training and validation action.\n",
    "* The shape of datasets is printed in the console.\n",
    "* All these actions are performed using the `return_dataloader()` and `CustomDataset class` defined above.\n",
    "\n",
    "\n",
    "<a id='section504'></a>\n",
    "#### Neural Network and Optimizer\n",
    "\n",
    "* In this stage we define the model and optimizer that will be used for training and to update the weights of the network. \n",
    "* We are using the `roberta-base-uncased` transformer model for our project. You can read about the `RoBERTa model` and its features above. \n",
    "* The model is returned and instiated using the `return_model()` and `ModelClass`.\n",
    "* We are using the `Adam` optimizer for our project. This has been a standard for all our tutorials and is something that can be changed updated to see how different optimizer perform with different learning rates. \n",
    "* There is also a scope for doing more with Optimizer such a decay, momentum to dynamically update the Learning rate and other parameters. All those concepts have been kept out of scope for these tutorials. \n",
    "\n",
    "\n",
    "<a id='section505'></a>\n",
    "#### Training Model and Logging to WandB\n",
    "\n",
    "* Followed by that we call the `train()` with all the necessary parameters.\n",
    "* Loss and accuracy at every 100th step is logged to the WandB service. \n",
    "* Accuracy and end of every epoch is logged in WandB and also printed in the console.\n",
    "\n",
    "\n",
    "<a id='section506'></a>\n",
    "#### Validation\n",
    "\n",
    "* After the training is completed, the validation step is initiated.\n",
    "* As defined in the validation function, the model weights are not updated. We use the fine tuned model to generate encoded sentiment.\n",
    "* An output is printed on the console giving the accuracy at the end of Validation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zgZVNTmu8w0J"
   },
   "outputs": [],
   "source": [
    "def run():\n",
    "    \n",
    "    # WandB – Initialize a new run\n",
    "    wandb.init(project=\"transformers_tutorials_sentiment\")\n",
    "    \n",
    "    # Defining some key variables that will be used later on in the training\n",
    "    config = wandb.config \n",
    "    config.MAX_LEN = 512\n",
    "    config.TRAIN_BATCH_SIZE = 4\n",
    "    config.VALID_BATCH_SIZE = 2\n",
    "    config.EPOCHS = 2\n",
    "    config.LEARNING_RATE = 1e-05\n",
    "    tokenizer = RobertaTokenizerFast.from_pretrained('roberta-base')\n",
    "\n",
    "    # Reading the dataset and pre-processing it for usage\n",
    "    df = pd.read_csv('/content/drive/My Drive/Colab Notebooks/IMDB Dataset.csv', encoding='latin-1')\n",
    "    pre = Preprocess(df)\n",
    "    encoding_dict, df = pre.processing()\n",
    "\n",
    "    # Creating the training and validation dataloader using the functions defined above\n",
    "    training_loader, validation_loader = return_dataloader(df, tokenizer, config.TRAIN_BATCH_SIZE, config.VALID_BATCH_SIZE, config.MAX_LEN)\n",
    "\n",
    "    # Defining the model based on the function and ModelClass defined above\n",
    "    model = return_model(device)\n",
    "\n",
    "    # Creating the loss function and optimizer\n",
    "    loss_function = torch.nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.Adam(params =  model.parameters(), lr=config.LEARNING_RATE)\n",
    "\n",
    "    # Fine tuning the model using the train function:\n",
    "    for epoch in range(config.EPOCHS):\n",
    "        train(epoch, model, device, training_loader, optimizer, loss_function)\n",
    "\n",
    "    # Running the validation function to validate the performance of the trained model\n",
    "    valid(epoch, model, device, validation_loader, loss_function)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 155
    },
    "colab_type": "code",
    "id": "MzQHSMj7IDIz",
    "outputId": "005fedb8-9be9-4809-d2e4-661823b98021",
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                Logging results to <a href=\"https://wandb.com\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
       "                Project page: <a href=\"https://app.wandb.ai/abhimishra-91/transformers_tutorials_sentiment\" target=\"_blank\">https://app.wandb.ai/abhimishra-91/transformers_tutorials_sentiment</a><br/>\n",
       "                Run page: <a href=\"https://app.wandb.ai/abhimishra-91/transformers_tutorials_sentiment/runs/1zwn4gbg\" target=\"_blank\">https://app.wandb.ai/abhimishra-91/transformers_tutorials_sentiment/runs/1zwn4gbg</a><br/>\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {
      "tags": []
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FULL Dataset: (50000, 2)\n",
      "TRAIN Dataset: (35000, 2)\n",
      "VAL Dataset: (15000, 2)\n",
      "The Total Accuracy for Epoch 0: 91.74285714285715\n",
      "The Total Accuracy for Epoch 1: 95.54\n",
      "The Validation Accuracy: 94.68\n"
     ]
    }
   ],
   "source": [
    "run()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "TPU",
  "colab": {
   "collapsed_sections": [],
   "name": "transformers_sentiment_imdb.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}